import asyncio
import datetime
import hashlib
import hmac
import json
import base64  # 添加这一行
from urllib.parse import urlparse, urlencode
from wsgiref.handlers import format_date_time
from time import mktime
import ssl
import websockets

history = []
# 用于存储当前连接的客户端
current_client = None

class Ws_Param(object):
    # 初始化
    def __init__(self, APPID, APIKey, APISecret, gpt_url):
        self.APPID = APPID
        self.APIKey = APIKey
        self.APISecret = APISecret
        self.host = urlparse(gpt_url).netloc
        self.path = urlparse(gpt_url).path
        self.gpt_url = gpt_url

    # 生成url
    def create_url(self):
        # 生成RFC1123格式的时间戳
        now = datetime.datetime.now()
        date = format_date_time(mktime(now.timetuple()))

        # 拼接字符串
        signature_origin = "host: " + self.host + "\n"
        signature_origin += "date: " + date + "\n"
        signature_origin += "GET " + self.path + " HTTP/1.1"

        # 进行hmac-sha256进行加密
        signature_sha = hmac.new(self.APISecret.encode('utf-8'), signature_origin.encode('utf-8'),
                                 digestmod=hashlib.sha256).digest()
        signature_sha_base64 = base64.b64encode(signature_sha).decode(encoding='utf-8')
        authorization_origin = f'api_key="{self.APIKey}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha_base64}"'
        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode(encoding='utf-8')
        # 将请求的鉴权参数组合为字典
        v = {
            "authorization": authorization,
            "date": date,
            "host": self.host
        }
        # 拼接鉴权参数，生成url
        url = self.gpt_url + '?' + urlencode(v)
        return url

async def connect_to_gpt(appid, api_secret, api_key, gpt_url, domain, query, history):
    wsParam = Ws_Param(appid, api_key, api_secret, gpt_url)
    wsUrl = wsParam.create_url()
    async with websockets.connect(wsUrl) as ws:
        data = json.dumps(gen_params(appid=appid, query=query, domain=domain, history=history))
        await ws.send(data)
        async for message in ws:
            data = json.loads(message)
            code = data['header']['code']
            if code != 0:
                print(f'请求错误: {code}, {data}')
                break
            else:
                choices = data["payload"]["choices"]
                status = choices["status"]
                content = choices["text"][0]["content"]
                print(content, end='')
                # 发送消息到客户端
                if current_client:
                    await current_client.send(content)
                # 将模型的回复添加到历史对话中
                history.append({"role": "assistant", "content": content})
                if status == 2:
                    print("\r\n")
                    break

def gen_params(appid, query, domain, history):
    """
    通过appid和用户的提问来生成请参数
    """
    data = {
        "header": {
            "app_id": appid,
            "uid": "1234",
        },
        "parameter": {
            "chat": {
                "domain": domain,
                "temperature": 0.5,
                "max_tokens": 4096,
                "auditing": "default",
            }
        },
        "payload": {
            "message": {
                "text": history + [{"role": "user", "content": query}]
            }
        }
    }
    return data

async def handle_connection(websocket, path):
    global current_client
    if current_client is not None:
        # 如果已有客户端连接，拒绝新的连接
        await websocket.send("服务器已有客户端连接，拒绝新连接")
        return

    # 建立新的客户端连接
    current_client = websocket
    try:
        print("客户端已连接")
        while True:
            # 接收客户端发送的消息
            message = await websocket.recv()
            print(f"收到客户端消息: {message}")
            history.append({"role": "user", "content": message})
            await connect_to_gpt(
                appid="8dbb4b5e",
                api_secret="NzZiZDUwMmFmYzIwZTRiMjVhNzQ0NjNk",
                api_key="e51a7fb80f7a69d365dcecb6b9a01dc6",
                gpt_url="wss://spark-api.xf-yun.com/v3.5/chat",
                domain="generalv3.5",
                query=message,
                history=history
            )
    except websockets.exceptions.ConnectionClosedOK:
        pass
    finally:
        # 当连接关闭时，清除客户端连接
        current_client = None
        print("客户端连接已关闭")

async def start_server():
    server = await websockets.serve(handle_connection, "192.168.8.149", 8765)
    print("服务器已启动")
    await server.wait_closed()

if __name__ == "__main__":
    asyncio.run(start_server())